# coding:UTF-8
# Version: V1.0.1
import serial
import time
import csv
import os

# --- 全局变量 ---
ACCData = [0.0]*8
GYROData = [0.0]*8
AngleData = [0.0]*8
FrameState = 0  # 帧状态判断
Bytenum = 0  # 当前帧已接收字节数
CheckSum = 0  # 校验和

a = [0.0]*3  # 加速度
w = [0.0]*3  # 角速度
Angle = [0.0]*3  # 角度

# --- 数据处理函数 ---
def DueData(inputdata):  # 核心数据处理函数
    global FrameState
    global Bytenum
    global CheckSum
    global a
    global w
    global Angle
    
    for data in inputdata:  # 遍历输入数据
        if FrameState == 0:  # 状态0：寻找帧头
            if data == 0x55 and Bytenum == 0:  # 找到第一个帧头字节0x55
                CheckSum = data
                Bytenum = 1
                continue
            elif data == 0x51 and Bytenum == 1:  # 加速度数据帧
                CheckSum += data
                FrameState = 1
                Bytenum = 2
            elif data == 0x52 and Bytenum == 1:  # 角速度数据帧
                CheckSum += data
                FrameState = 2
                Bytenum = 2
            elif data == 0x53 and Bytenum == 1:  # 角度数据帧
                CheckSum += data
                FrameState = 3
                Bytenum = 2
        elif FrameState == 1:  # 状态1：接收加速度数据
            if Bytenum < 10:  # 读取8个数据字节
                ACCData[Bytenum-2] = data
                CheckSum += data
                Bytenum += 1
            else:
                if data == (CheckSum & 0xff):  # 校验和检查
                    a = get_acc(ACCData)
                CheckSum = 0
                Bytenum = 0
                FrameState = 0
        elif FrameState == 2:  # 状态2：接收角速度数据
            if Bytenum < 10:
                GYROData[Bytenum-2] = data
                CheckSum += data
                Bytenum += 1
            else:
                if data == (CheckSum & 0xff):
                    w = get_gyro(GYROData)
                CheckSum = 0
                Bytenum = 0
                FrameState = 0
        elif FrameState == 3:  # 状态3：接收角度数据
            if Bytenum < 10:
                AngleData[Bytenum-2] = data
                CheckSum += data
                Bytenum += 1
            else:
                if data == (CheckSum & 0xff):
                    Angle = get_angle(AngleData)
                    print("加速度: %10.3f %10.3f %10.3f" % a)
                    print("角速度: %10.3f %10.3f %10.3f" % w)
                    print("角度: %10.3f %10.3f %10.3f" % Angle)
                    print("-" * 50)
                    return True  # 返回True表示成功解析了一组完整数据
                CheckSum = 0
                Bytenum = 0
                FrameState = 0
    return False  # 返回False表示没有解析到完整数据

def get_acc(datahex):
    axl = datahex[0]
    axh = datahex[1]
    ayl = datahex[2]
    ayh = datahex[3]
    azl = datahex[4]
    azh = datahex[5]
    k_acc = 16.0
    acc_x = (axh << 8 | axl) / 32768.0 * k_acc
    acc_y = (ayh << 8 | ayl) / 32768.0 * k_acc
    acc_z = (azh << 8 | azl) / 32768.0 * k_acc
    if acc_x >= k_acc:
        acc_x -= 2 * k_acc
    if acc_y >= k_acc:
        acc_y -= 2 * k_acc
    if acc_z >= k_acc:
        acc_z -= 2 * k_acc
    return acc_x, acc_y, acc_z

def get_gyro(datahex):
    wxl = datahex[0]
    wxh = datahex[1]
    wyl = datahex[2]
    wyh = datahex[3]
    wzl = datahex[4]
    wzh = datahex[5]
    k_gyro = 2000.0
    gyro_x = (wxh << 8 | wxl) / 32768.0 * k_gyro
    gyro_y = (wyh << 8 | wyl) / 32768.0 * k_gyro
    gyro_z = (wzh << 8 | wzl) / 32768.0 * k_gyro
    if gyro_x >= k_gyro:
        gyro_x -= 2 * k_gyro
    if gyro_y >= k_gyro:
        gyro_y -= 2 * k_gyro
    if gyro_z >= k_gyro:
        gyro_z -= 2 * k_gyro
    return gyro_x, gyro_y, gyro_z

def get_angle(datahex):
    rxl = datahex[0]
    rxh = datahex[1]
    ryl = datahex[2]
    ryh = datahex[3]
    rzl = datahex[4]
    rzh = datahex[5]
    k_angle = 180.0
    angle_x = (rxh << 8 | rxl) / 32768.0 * k_angle
    angle_y = (ryh << 8 | ryl) / 32768.0 * k_angle
    angle_z = (rzh << 8 | rzl) / 32768.0 * k_angle
    if angle_x >= k_angle:
        angle_x -= 2 * k_angle
    if angle_y >= k_angle:
        angle_y -= 2 * k_angle
    if angle_z >= k_angle:
        angle_z -= 2 * k_angle
    return angle_x, angle_y, angle_z

if __name__ == '__main__':
    port = '/dev/ttyUSB0'  # IMU串口号
    baud = 9600  # 与IMU模块相同的波特率
    
    try:
        ser = serial.Serial(port, baud, timeout=0.5)
        print("串口已打开:", ser.is_open)
        print("正在读取IMU数据...")
        print("按Ctrl+C停止程序")
        
        # 获取用户输入的文件名
        drive_command = input("请输入文件名（数字）: ")
        filename = f"{drive_command}.csv"
        
        # 创建CSV文件并写入表头
        with open(filename, 'w', newline='') as csvfile:
            csv_writer = csv.writer(csvfile)
            csv_writer.writerow(['Timestamp', 'Acc_X', 'Acc_Y', 'Acc_Z', 
                               'Gyro_X', 'Gyro_Y', 'Gyro_Z',
                               'Angle_X', 'Angle_Y', 'Angle_Z'])
            
            start_time = time.time()
            data_count = 0
            
            while True:
                datahex = ser.read(33)  # 读取33字节数据
                if datahex:
                    if DueData(datahex):  # 如果成功解析了一组数据
                        timestamp = time.time() - start_time
                        # 写入数据到CSV文件
                        csv_writer.writerow([
                            timestamp,
                            a[0], a[1], a[2],
                            w[0], w[1], w[2],
                            Angle[0], Angle[1], Angle[2]
                        ])
                        csvfile.flush()  # 确保数据立即写入文件
                        data_count += 1
                        
                        if data_count % 100 == 0:  # 每100个数据点显示一次进度
                            print(f"已记录 {data_count} 个数据点")
            
    except KeyboardInterrupt:
        print("\n程序已停止")
        print(f"共记录 {data_count} 个数据点")
    except serial.SerialException as e:
        print(f"串口错误: {e}")
    finally:
        if 'ser' in locals() and ser.is_open:
            ser.close()
            print("串口已关闭") 